package com.lin.linchatgpt.listener;

import cn.hutool.core.util.StrUtil;
import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.lin.linchatgpt.entity.Message;
import com.lin.linchatgpt.exception.LinException;
import com.lin.linchatgpt.service.MessageService;
import com.lin.linchatgpt.utils.NanoMillisUtils;
import com.lin.linchatgpt.utils.ResultCode;
import com.unfbx.chatgpt.entity.chat.ChatChoice;
import com.unfbx.chatgpt.entity.chat.ChatCompletionResponse;
import lombok.NoArgsConstructor;
import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;
import okhttp3.Response;
import okhttp3.ResponseBody;
import okhttp3.sse.EventSource;
import okhttp3.sse.EventSourceListener;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import org.springframework.transaction.annotation.Transactional;
import org.springframework.web.servlet.mvc.method.annotation.ResponseBodyEmitter;

import javax.annotation.Resource;
import java.io.IOException;
import java.util.List;
import java.util.Objects;
import java.util.Timer;
import java.util.concurrent.ThreadPoolExecutor;

/**
 * Copyright(C),2022年-2022,霖霖
 *
 * @author 霖霖
 * @version 1.0
 * @date 2023/6/3 11:31
 * @Description
 */
@Slf4j
public class ParsedEventSourceListener extends EventSourceListener {
    private final ResponseBodyEmitter emitter;
    private final StringBuilder answer = new StringBuilder();

    private final MessageService messageService;
    private final String subjectId;
    private final String parentMessageId;
    private final Message questionMessage;
    private final ThreadPoolExecutor taskExecutor;
    private int type;

    private ParsedEventSourceListener(Builder builder) {
        this.messageService = builder.messageService;
        this.subjectId = builder.subjectId;
        this.parentMessageId = builder.parentMessageId;
        this.emitter = builder.emitter;
        this.questionMessage = builder.questionMessage;
        this.taskExecutor = builder.taskExecutor;
        this.type = builder.type;
    }

    @Override
    public void onOpen(@NotNull EventSource eventSource, @NotNull Response response) {
        // 设置定时器，如果超过60秒没有返回数据，则关闭连接
        log.info("OpenAI连接开启,{}", response);
    }

    @Override
    public void onEvent(@NotNull EventSource eventSource, @Nullable String id, @Nullable String type, @NotNull String data) {

        taskExecutor.execute(()->{
            if (data.equals("[DONE]")) {
                log.info("OpenAI返回数据结束");
                emitter.complete();

                Message message = new Message();
                message.setContent(answer.toString());
                message.setSubjectId(subjectId);
                message.setRole(com.unfbx.chatgpt.entity.chat.Message.Role.ASSISTANT.getName());
                message.setParentMessageId(parentMessageId);
                message.setRecordTime(NanoMillisUtils.getMillisTime());
                try {
                    messageService.saveMessage(questionMessage, message, this.type);
                } catch (Exception e) {
                    e.printStackTrace();
                }
                return;
            }
            ChatCompletionResponse chatCompletionResponse = null;
            try {
                chatCompletionResponse = new ObjectMapper().readValue(data, ChatCompletionResponse.class);
            } catch (JsonProcessingException e) {
                e.printStackTrace();
                throw new LinException(ResultCode.OPENAI_SSE_ERROR, "OpenAI返回数据解析失败");
            }
            List<ChatChoice> choices = chatCompletionResponse.getChoices();
            String content = choices.get(0).getDelta().getContent();
            if (StrUtil.isNotEmpty(content)) {
                answer.append(content);
                try {
                    emitter.send(content);
                } catch (IOException e) {
                    e.printStackTrace();
                    throw new LinException(ResultCode.OPENAI_SSE_ERROR, "OpenAI返回数据发送给用户失败");
                }
            }
        });
    }

    @SneakyThrows
    @Override
    public void onFailure(@NotNull EventSource eventSource, @Nullable Throwable t, @Nullable Response response) {
        String responseError = "";
        ResponseBody body = null;
        int resCode = 0;

        if (Objects.nonNull(response)) {
            body = response.body();
            responseError = response.toString();
            log.error("OpenAI连接失败response,{},状态码：{}", response, response.code());
            resCode = response.code();
        }
        if (Objects.nonNull(body)) {
            responseError = body.string();
            log.error("OpenAI连接失败body,{}", responseError);
        }
        if (t != null) {
            responseError = t.getMessage();
            log.error("OpenAI连接失败t,{}", responseError);
        }
        log.error("OpenAI连接失败,{}", responseError);
        emitter.completeWithError(new LinException(resCode != 0 ? resCode : ResultCode.OPENAI_SSE_ERROR, responseError));

        eventSource.cancel();
    }

    @Override
    public void onClosed(@NotNull EventSource eventSource) {
        log.info("OpenAI连接关闭...");
    }


    public static class Builder {
        private ResponseBodyEmitter emitter;
        private String subjectId;
        private String parentMessageId;
        private MessageService messageService;
        private Message questionMessage;

        private ThreadPoolExecutor taskExecutor;
        private int type;

        public Builder setEmitter(ResponseBodyEmitter emitter) {
            this.emitter = emitter;
            return this;
        }

        public Builder setSubjectId(String subjectId) {
            this.subjectId = subjectId;
            return this;
        }

        public Builder setParentMessageId(String parentMessageId) {
            this.parentMessageId = parentMessageId;
            return this;
        }

        public Builder setMessageService(MessageService messageService) {
            this.messageService = messageService;
            return this;
        }

        public Builder setQuestionMessage(Message questionMessage) {
            this.questionMessage = questionMessage;
            return this;
        }

        public Builder setTaskExecutor(ThreadPoolExecutor taskExecutor) {
            this.taskExecutor = taskExecutor;
            return this;
        }

        public Builder setType(int type) {
            this.type = type;
            return this;
        }

        public ParsedEventSourceListener build() {
            if (Objects.isNull(emitter)) {
                throw new IllegalStateException("emitter is null");
            }
            return new ParsedEventSourceListener(this);
        }
    }
}
